#!/usr/bin/env python
import os
import socket
import subprocess
import sys
from pathlib import Path

import numpy as np
import setproctitle
import torch

import wandb

script_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(script_dir)))
sys.path.append(parent_dir)
from onpolicy.config_env import get_env_config
from onpolicy.config_train import get_config
from onpolicy.envs.env_continuous import ContinuousActionEnv
from onpolicy.envs.env_wrappers import DummyVecEnv, SubprocVecEnv
from onpolicy.envs.mpe.MPE_env import MPEEnv


def get_gpu_memory_map():
    """Get the current gpu usage.

    Returns
    -------
    usage: dict
        Keys are device ids as integers.
        Values are memory usage as integers in MB.
    """
    result = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
        encoding="utf-8",
    )
    # Convert lines into a dictionary
    gpu_memory = [int(x) for x in result.strip().split("\n")]
    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))
    return gpu_memory_map


gpu_memory_map = get_gpu_memory_map()
min_memory_gpu = min(gpu_memory_map, key=gpu_memory_map.get)
Device = torch.device(f"cuda:{min_memory_gpu}" if torch.cuda.is_available() else "cpu")
os.environ["CUDA_VISIBLE_DEVICES"] = str(min_memory_gpu)
torch.cuda.set_device(Device)


def make_render_env(all_args, env_size=[]):
    def get_env_fn(rank):
        def init_env(env_size):
            if all_args.env_name == "MPE_real":
                env = MPEEnv(all_args)
            elif all_args.env_name == "MyEnv":
                env = ContinuousActionEnv(env_size)
            else:
                print("Can not support the " + all_args.env_name + "environment.")
                raise NotImplementedError
            env.seed(all_args.seed * 50000 + rank * 10000)
            return env

        return init_env

    if all_args.n_eval_rollout_threads == 1:
        return DummyVecEnv([get_env_fn(0)], env_size)
    else:
        return SubprocVecEnv(
            [get_env_fn(i) for i in range(all_args.n_eval_rollout_threads)], env_size
        )


def parse_args(args, parser):
    parser.add_argument(
        "--scenario_name",
        type=str,
        default="simple_spread",
        help="Which scenario to run on",
    )
    parser.add_argument("--num_agents", type=int, default=2, help="number of players")

    all_args = parser.parse_known_args(args)[0]

    return all_args


def main(args, env_size=[], device=Device):
    parser = get_config()
    all_args = parse_args(args, parser)

    if all_args.algorithm_name == "rmappo":
        print("u are choosing to use rmappo, we set use_recurrent_policy to be True")
        all_args.use_recurrent_policy = True
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "mappo":
        print(
            "u are choosing to use mappo, we set use_recurrent_policy & use_naive_recurrent_policy to be False"
        )
        all_args.use_recurrent_policy = False
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "ippo":
        print("u are choosing to use ippo, we set use_centralized_V to be False.")
        all_args.use_centralized_V = False
    else:
        raise NotImplementedError

    assert (
        (
            all_args.share_policy == True
            and all_args.scenario_name == "simple_speaker_listener"
        )
        == False
    ), "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    assert all_args.use_render, "u need to set use_render be True"
    assert not (
        all_args.model_dir == None or all_args.model_dir == ""
    ), "set model_dir first"
    assert all_args.n_rollout_threads == 1, "only support to use 1 env to render."

    # # cuda
    # if all_args.cuda and torch.cuda.is_available():
    #     print("choose to use gpu...")
    #     device = torch.device("cuda:0")
    #     torch.set_num_threads(all_args.n_training_threads)
    #     if all_args.cuda_deterministic:
    #         torch.backends.cudnn.benchmark = False
    #         torch.backends.cudnn.deterministic = True
    # else:
    #     print("choose to use cpu...")
    #     device = torch.device("cpu")
    #     torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    # wandb
    # if all_args.use_wandb:
    #     run = wandb.init(
    #         config=all_args,
    #         project=all_args.env_name,
    #         # entity=all_args.user_name,
    #         notes=socket.gethostname(),
    #         name=str(all_args.algorithm_name) + "_" + "render",
    #         # group=all_args.scenario_name,
    #         dir=str(run_dir),
    #         # job_type="training",
    #         reinit=True,
    #     )
    # else:
    #     if not run_dir.exists():
    #         curr_run = "run1"
    #     else:
    #         exst_run_nums = [
    #             int(str(folder.name).split("run")[1])
    #             for folder in run_dir.iterdir()
    #             if str(folder.name).startswith("run")
    #         ]
    #         if len(exst_run_nums) == 0:
    #             curr_run = "run1"
    #         else:
    #             curr_run = "run%i" % (max(exst_run_nums) + 1)
    #     run_dir = run_dir / curr_run
    #     if not run_dir.exists():
    #         os.makedirs(str(run_dir))

    # setproctitle.setproctitle(
    #     str(all_args.algorithm_name)
    #     + "-"
    #     + str(all_args.env_name)
    #     + "-"
    #     + str(all_args.experiment_name)
    #     + "@"
    #     + str(all_args.user_name)
    # )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_render_env(all_args, env_size)
    eval_envs = None
    num_agents = all_args.num_agents

    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from onpolicy.runner.shared.mpe_runner import MPERunner as Runner
    else:
        from onpolicy.runner.separated.mpe_runner import MPERunner as Runner

    runner = Runner(config)
    runner.render()

    # post process
    envs.close()

    # if all_args.use_wandb:
    #     run.finish()


from onpolicy.config_env import get_env_config
from onpolicy.config_train import get_config

# from tool.email import Mail
from onpolicy.envs.env_wrappers import DummyVecEnv
from onpolicy.envs.ruler import ruler_main, show_move_list

parser_ev = get_env_config()

if __name__ == "__main__":
    env = "MyEnv"
    scenario = "simple_spread"

    num_agents = 3
    algo = "rmappo"
    exp = "check"
    seed_max = 1
    model_num = sys.argv[1]
    # 设置 sys.argv 的值
    argv = [
        "--env_name",
        f"{env}",
        "--algorithm_name",
        f"{algo}",
        "--experiment_name",
        f"{exp}",
        "--scenario_name",
        f"{scenario}",
        "--num_agents",
        f"{num_agents}",
        "--seed",
        f"{seed_max}",
        "--n_training_threads",
        "1",
        "--n_rollout_threads",
        "1",
        "--num_mini_batch",
        "1",
        "--episode_length",
        "200",
        "--num_env_steps",
        "20000000",
        "--ppo_epoch",
        "20",
        "--gain",
        "0.01",
        "--lr",
        "1e-3",
        "--critic_lr",
        "1e-5",
        "--model_dir",
        "E:/Code_file/Python/Env_1106/onpolicy/scripts/results/MyEnv/simple_spread/rmappo/check/run"
        + model_num
        + "/models",
        "--wandb_name",
        "--use_render",
        "test_wandb_mpe",
    ]
    sys.argv = argv

    # if parser_ev["is_Draw_now"] and parser_ev["env_type"] == "Draw":
    #     ruler_main()
    #     show_move_list()
    # elif parser_ev["env_type"] == "Rand":
    #     for i in range(parser_ev["rand_type"]):
    #         for j in range(parser_ev["rand_epo"]):
    #             main(sys.argv[1:], ["Circle", [-i, i + 1], 10 + i], Device)
    #     for i in range(parser_ev["rand_type"]):
    #         for j in range(parser_ev["rand_epo"]):
    #             main(
    #                 sys.argv,
    #                 ["Rectangle", [-i - 10, i - 10], 10 + i, 20 - i],
    #                 Device,
    #             )
    #     for i in range(parser_ev["rand_type"]):
    #         for j in range(parser_ev["rand_epo"]):
    #             main(sys.argv[1:], ["Circle", [-i + 1, i + 1], 5 + i], Device)
    #     for i in range(parser_ev["rand_type"]):
    #         for j in range(parser_ev["rand_epo"]):
    #             main(
    #                 sys.argv[1:],
    #                 ["Rectangle", [-20 + i, 20 - i * 2], 5 + i * 2, 20 - i],
    #                 Device,
    #             )
    # else:
    main(sys.argv[1:], ["Circle", [0, 1], 10], device=Device)
    main(sys.argv[1:], ["Rectangle", [10, 10], 10, 10], device=Device)
